import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules
import pickle
import matplotlib.pyplot as plt
import networkx as nx
from flask import Flask, request, jsonify
from itertools import combinations

# 设置中文字体显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# -------------------- 数据处理与关联规则挖掘 --------------------
def main():
    # 读取数据
    df = pd.read_excel('D:\作业\专业课\数据挖掘\作业2\关联规则\餐厅数据.xlsx')
    print("数据前五行预览：")
    print(df.head())

    # 数据预处理 - 将菜品列转换为事务列表
    transactions = df['菜品'].str.split(',').tolist()

    # 编码事务数据
    te = TransactionEncoder()
    te_ary = te.fit(transactions).transform(transactions)
    df_encoded = pd.DataFrame(te_ary, columns=te.columns_)

    # 挖掘频繁项集
    frequent_itemsets = apriori(df_encoded, min_support=0.1, use_colnames=True)
    frequent_itemsets.sort_values(by='support', ascending=False, inplace=True)
    
    # 打印二元频繁项集
    print("\n二元频繁项集：")
    print(frequent_itemsets[frequent_itemsets['itemsets'].apply(lambda x: len(x) == 2)])

    # 生成关联规则
    rules = association_rules(frequent_itemsets, metric="confidence", min_threshold=0.1)
    
    # 筛选有效规则（提升度>1且置信度>0.1）
    effective_rules = rules[
        (rules['lift'] > 1) & (rules['confidence'] > 0.1)
    ].sort_values(by=['lift', 'confidence'], ascending=False)
    
    # 打印有效规则
    print("\n有效关联规则：")
    print(effective_rules[['antecedents', 'consequents', 'support', 'confidence', 'lift']])

    # 保存模型
    with open('frequent_itemsets.pkl', 'wb') as f:
        pickle.dump(frequent_itemsets, f)
        
    with open('rules.pkl', 'wb') as f:
        pickle.dump(effective_rules, f)

    # 可视化关联规则
    visualize_rules(effective_rules)

# -------------------- 关联规则可视化 --------------------
def visualize_rules(rules):
    G = nx.DiGraph()
    
    # 添加节点和边（规则）
    for _, row in rules.iterrows():
        antecedents = ','.join(list(row['antecedents']))
        consequents = ','.join(list(row['consequents']))
        G.add_edge(antecedents, consequents, weight=row['lift'])
    
    # 设置图形
    plt.figure(figsize=(12, 8))
    pos = nx.spring_layout(G)
    
    # 绘制网络图
    nx.draw(G, pos,
            with_labels=True,
            edge_color=[G[u][v]['weight'] for u, v in G.edges()],
            width=2.0,
            edge_cmap=plt.cm.Blues)
    
    plt.title("关联规则网络图")
    plt.show()

# -------------------- 推荐API服务 --------------------
app = Flask(__name__)

# 加载模型
frequent_itemsets = pd.read_pickle('./frequent_itemsets.pkl')
rules = pd.read_pickle('./rules.pkl')

@app.route('/recommend', methods=['POST'])
def recommend():
    # 获取用户已选菜品
    selected_items = request.json.get('items', [])
    
    # 基于关联规则生成推荐
    recommendations = []
    for _, rule in rules.iterrows():
        antecedents = list(rule['antecedents'])
        consequents = list(rule['consequents'])
        
        # 如果规则前件包含用户已选菜品，则推荐后件
        if set(antecedents).issubset(set(selected_items)):
            recommendations.extend(consequents)
    
    # 去重并排除已选菜品
    recommendations = list(set(recommendations) - set(selected_items))
    
    # 返回推荐结果
    return jsonify({"recommendations": recommendations})

if __name__ == '__main__':
    # 运行数据处理和规则挖掘
    main()
    
    # 启动API服务
    # 注意：生产环境建议设置debug=False
    app.run(debug=True)